import torch
from typing import List, Dict
from collections import OrderedDict
from tqdm import tqdm

from hj_dataset_devkit import Frame
from model_frame import build_task_assigner, build_model, HJItem, DetectObject

from common_utils import read_yaml
from components.obj_detector import ObjDetector

__all__ = ['VoxelNeXtDetector']

def load_state_to_model(pth_path: str, model: torch.nn.Module) -> None:
    def has_module_prefix(state: Dict[str, torch.Tensor]) -> bool:
        for key in state.keys():
            if key.startswith('module.'):
                return True
        return False
    
    def fix_model_state(in_state: Dict[str, torch.Tensor]) -> OrderedDict:
        fixed_state = OrderedDict()
        for k, v in in_state.items():
            fixed_state[k[7:]] = v
        return fixed_state
    
    checkpoint: dict = torch.load(pth_path, 'cpu')
    model_state: dict = checkpoint.get('model_state_dict')
    if has_module_prefix(model_state):
        model_state = fix_model_state(model_state)
    model.load_state_dict(model_state)

class VoxelNeXtDetector(ObjDetector):
    def __init__(self, train_cfg_path: str, model_pth_path: str, device: str) -> None:
        super().__init__()
        train_cfg = read_yaml(train_cfg_path)
        self.__task = build_task_assigner(train_cfg['task'])
        self.__model = build_model(train_cfg['model'])
        load_state_to_model(model_pth_path, self.__model)
        print(f'INFO: detector model is loaded from {model_pth_path}')
        self.__device = torch.device(device)
        self.__model.to(self.__device)
        self.__model.eval()

    def process_frame(self, frame: Frame) -> List[DetectObject]:
        item = HJItem(frame.get_lidar_cloud('true_value')[:, :4], [])
        in_data = self.__task.assign_input(item)
        in_data, _ = self.__task.collate_fn(
                [(in_data, self.__task.assign_target(item))])
        with torch.no_grad():
            out_data = self.__model(in_data.to(self.__device))
        return self.__task.parse_model_output(out_data)[0]

    def process_frames_by_batch(self, frames: List[Frame], batch_size=8) -> Dict[int, List[DetectObject]]:
        result: Dict[int, List[DetectObject]] = {}
        pbar = tqdm(total=len(frames), desc=f'detect by batch {batch_size}')
        for i in range(len(frames) // batch_size + 1):
            batch_frames = frames[i*batch_size: min((i+1)*batch_size, len(frames))]
            if len(batch_frames) == 0:
                continue
            items = [HJItem(frame.get_lidar_cloud('true_value')[:, :4], [])
                     for frame in batch_frames]
            in_data, _ = self.__task.collate_fn([
                (self.__task.assign_input(it), self.__task.assign_target(it)) for it in items
            ])
            with torch.no_grad():
                out_data = self.__model(in_data.to(self.__device))
            parsed_output = self.__task.parse_model_output(out_data)
            for j, frame in enumerate(batch_frames):
                result[frame.get_timestamp()] = parsed_output[j]
            pbar.update(len(batch_frames))
        return result
